from fab.trainable_distributions.base import TrainableDistribution